package com.jd.genie.agent.tool.common;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.jd.genie.agent.agent.AgentContext;
import com.jd.genie.agent.dto.CodeInterpreterRequest;
import com.jd.genie.agent.dto.CodeInterpreterResponse;
import com.jd.genie.agent.dto.File;
import com.jd.genie.agent.printer.Printer;
import com.jd.genie.agent.tool.BaseTool;
import com.jd.genie.agent.util.SpringContextHolder;
import com.jd.genie.config.GenieConfig;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.context.ApplicationContext;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

@Slf4j
public class CodeInterpreterTool implements BaseTool {

    private final AgentContext agentContext;

    private final Printer printer;

    public CodeInterpreterTool(AgentContext agentContext, Printer printer)
    {
        this.agentContext = agentContext;
        this.printer = printer;
    }

    @Override
    public String getName()
    {
        return "code_interpreter";
    }

    @Override
    public String getDescription()
    {
        String desc = "这是一个代码工具，可以通过编写代码完成数据处理、数据分析、图表生成等任务";
        return """
                这是一个Code interpreter工具，可以写Python代码
                - 严禁用此工具进行处理从非表格文件中提取表格、抽取数据、抽取指标等任务。
                - 严禁处理纯文本文件，例如 .txt, .md , .html 等文件的直接处理。如果需要对这些文件分析，则应该先通过读取这些文件内容，保存成 .csv 文件格式的数据表后进行处理。
                - 如果上下文中有.xlsx 、.csv 等Excel表格文件需要处理分析，可以直接使用此工具读取 .xlsx 、.csv 等Excel表格文件进行分析处理
                
                """;
    }

    /**
     * codeAgent 配置
     */
    private final Map<String, Object> codeAgentParams;

    {
        String jsonStr = """
                {
                  "type" : "object",
                  "properties" : {
                    "task" : {
                      "description" : "任务的描述，不仅包括提供任务目标、任务的相关要求，还包括详细的完成任务需要的细节信息。详细是指不仅包含上下文中提及的所有与任务的相关内容，同时，包括：用户提供的业务名词、业务背景、数据等相关内容，确保这是一个易于理解、步骤明确、且能完成完整的任务描述。完整任务的定义是所有依赖项都在这里写清楚，明确无歧义，信息无丢失，基于这些信息足够完成该任务。禁止编造数据，写出来的程序是基于上下文已有的数据进行。",
                      "type" : "string"
                    }
                  },
                  "required" : [ "task" ]
                }
                """;
        codeAgentParams = JSON.parseObject(jsonStr, Map.class);
    }

    @Override
    public Map<String, Object> toParams()
    {
        return codeAgentParams;
    }

    @Override
    public Object execute(Object input)
    {
        try {
            Map<String, Object> params = (Map<String, Object>) input;
            String task = (String) params.get("task");
            List<String> fileNames = agentContext.getProductFiles().stream().map(File::getFileName).collect(Collectors.toList());
            CodeInterpreterRequest request = CodeInterpreterRequest.builder()
                    .requestId(agentContext.getSessionId()) // 适配多轮对话
                    .query(agentContext.getQuery())
                    .task(task)
                    .fileNames(fileNames)
                    .stream(true)
                    .build();

            return callCodeAgentStream(request);
        } catch (Exception e) {
            log.error("{} code agent error", agentContext.getRequestId(), e);
        }
        return null;
    }

    /**
     * 调用 CodeAgent
     */
    @SneakyThrows
    public String callCodeAgentStream(CodeInterpreterRequest codeRequest)
    {
        OkHttpClient client = new OkHttpClient.Builder()
                .connectTimeout(60, TimeUnit.SECONDS) // 设置连接超时时间为 60 秒
                .readTimeout(300, TimeUnit.SECONDS)    // 设置读取超时时间为 60 秒
                .writeTimeout(300, TimeUnit.SECONDS)   // 设置写入超时时间为 60 秒
                .callTimeout(300, TimeUnit.SECONDS)    // 设置调用超时时间为 60 秒
                .build();

        ApplicationContext applicationContext = SpringContextHolder.getApplicationContext();
        GenieConfig genieConfig = applicationContext.getBean(GenieConfig.class);
        String url = genieConfig.getCodeInterpreterUrl() + "/v1/tool/code_interpreter";
        RequestBody body = RequestBody.create(
                MediaType.parse("application/json"),
                JSONObject.toJSONString(codeRequest)
        );

        log.info("{} code_interpreter request {}", agentContext.getRequestId(), JSONObject.toJSONString(codeRequest));
        Request.Builder requestBuilder = new Request.Builder()
                .url(url)
                .post(body);
        Request request = requestBuilder.build();

        CodeInterpreterResponse codeResponse = CodeInterpreterResponse.builder()
                .codeOutput("code_interpreter执行失败") // 默认输出
                .build();
        try (Response response = client.newCall(request).execute();
             ResponseBody responseBody = response.body()) {
            log.info("{} code_interpreter response {} {} {}", agentContext.getRequestId(), response, response.code(), response.body());
            if (!response.isSuccessful() || responseBody == null) {
                log.error("{} code_interpreter request error", agentContext.getRequestId());
                throw new IOException("Unexpected response code: " + response);
            }

            String line;
            BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream()));
            while ((line = reader.readLine()) != null) {
                if (line.startsWith("data: ")) {
                    String data = line.substring(6);
                    if (data.equals("[DONE]")) {
                        break;
                    }
                    if (data.startsWith("heartbeat")) {
                        continue;
                    }
                    log.info("{} code_interpreter recv data: {}", agentContext.getRequestId(), data);
                    codeResponse = JSONObject.parseObject(data, CodeInterpreterResponse.class);
                    if (Objects.nonNull(codeResponse.getFileInfo()) && !codeResponse.getFileInfo().isEmpty()) {
                        for (CodeInterpreterResponse.FileInfo fileInfo : codeResponse.getFileInfo()) {
                            File file = File.builder()
                                    .fileName(fileInfo.getFileName())
                                    .ossUrl(fileInfo.getOssUrl())
                                    .domainUrl(fileInfo.getDomainUrl())
                                    .fileSize(fileInfo.getFileSize())
                                    .description(fileInfo.getFileName()) // fileName用作描述
                                    .isInternalFile(false)
                                    .build();
                            agentContext.getProductFiles().add(file);
                            agentContext.getTaskProductFiles().add(file);
                        }
                    }
                    String digitalEmployee = agentContext.getToolCollection().getDigitalEmployee(getName());
                    log.info("requestId:{} task:{} toolName:{} digitalEmployee:{}", agentContext.getRequestId(),
                            agentContext.getToolCollection().getCurrentTask(), getName(), digitalEmployee);
                    printer.send("code", codeResponse, digitalEmployee);
                }
            }

        }
        /**
         * {{输出内容}}
         * \n\n
         * 其中保存了文件：
         * {{文件名}}
         */
        StringBuilder output = new StringBuilder();
        output.append(codeResponse.getCodeOutput());
        if (Objects.nonNull(codeResponse.getFileInfo()) && !codeResponse.getFileInfo().isEmpty()) {
            output.append("\n\n其中保存了文件: ");
            for (CodeInterpreterResponse.FileInfo fileInfo : codeResponse.getFileInfo()) {
                output.append(fileInfo.getFileName()).append("\n");
            }
        }
        return output.toString();
    }

}